from common.kafka_message import get_factor_topic, get_kline_topic
from common.variables import time_counts, TimeFrameType
from datetime import datetime
from itertools import product
import json
from sortedcontainers import SortedDict
import time


class KafkaWatcherKline:
    def __init__(self, timeframe: TimeFrameType, watch_group: str):
        self.cur_time = None

        self.timeframe = timeframe
        self.last_time = time.time() - 2 * time_counts[timeframe]
        self.data = SortedDict()  # key: time, value: list[symbol]
        self.watch_group = watch_group

    def watch_batch_symbols(self, symbols: list[str]):
        from kafka import KafkaConsumer

        topics = [get_kline_topic(x, self.timeframe) for x in symbols]
        # 创建 Kafka 消费者
        consumer = KafkaConsumer(
            *topics,  # 订阅的 Kafka topic
            bootstrap_servers="xx_kafka:9092",  # Kafka 服务地址
            group_id=self.watch_group,  # 消费者组 ID
            auto_offset_reset="latest",
        )

        try:
            for message in consumer:
                try:
                    msg = message.value.decode("utf-8")
                    alltmp = json.loads(msg)
                    for tmp in alltmp:
                        t = tmp["time"] / 1000
                        if (
                            t + time_counts[self.timeframe] + 60 < time.time()
                        ):  # 当前时间比理论到达时间的60秒后还更晚
                            raise Exception(f"k线数据延迟严重, {datetime.now()}, {tmp}")
                        if t <= self.last_time:
                            continue
                        if t not in self.data:
                            self.data[t] = set()
                        self.data[t].add(tmp["symbol"])
                        if len(self.data[t]) >= len(symbols):
                            times = list(self.data.keys())
                            for xt in times:
                                if xt <= t:
                                    self.data.pop(xt)
                                    # if len(self.data[xt]) < len(symbols):  # 没有发出过信号的k线触发None
                                    #     yield None
                            self.last_time = t
                            yield t  # 满足要求的k线触发时间
                except Exception as e:
                    print(f"kafka消息错误{e}")
        finally:
            consumer.close()


class KafkaWatcherFactors:
    def __init__(self, timeframe: TimeFrameType, watch_group: str):
        self.cur_time = None

        self.timeframe = timeframe
        self.last_time = time.time() - 2 * time_counts[timeframe]
        self.data = SortedDict()  # key: time, value: list[symbol]
        self.watch_group = watch_group

    def watch_batch_symbol_factors(self, symbols: list[str], tfname_list: list[str]):
        from kafka import KafkaConsumer

        topics = [
            get_factor_topic(x, y, self.timeframe)
            for x, y in product(symbols, tfname_list)
        ]
        # 创建 Kafka 消费者
        consumer = KafkaConsumer(
            *topics,  # 订阅的 Kafka topic
            bootstrap_servers="xx_kafka:9092",  # Kafka 服务地址
            group_id=self.watch_group,  # 消费者组 ID
            auto_offset_reset="latest",
        )

        try:
            for message in consumer:
                try:
                    msg = message.value.decode("utf-8")
                    alltmp = json.loads(msg)
                    for tmp in alltmp:
                        t = tmp["time"] / 1000
                        if (
                            t + time_counts[self.timeframe] + 60 < time.time()
                        ):  # 当前时间比理论到达时间的60秒后还更晚
                            raise Exception(
                                f"因子数据延迟严重, {datetime.now()}, {tmp}"
                            )
                        if t <= self.last_time:
                            continue
                        if t not in self.data:
                            self.data[t] = set()
                        self.data[t].add((tmp["symbol"], tmp["tfname"]))
                        if len(self.data[t]) >= len(topics):
                            times = list(self.data.keys())
                            for xt in times:
                                if xt <= t:
                                    self.data.pop(xt)
                                    # if len(self.data[xt]) < len(symbols):  # 没有发出过信号的k线触发None
                                    #     yield None
                            self.last_time = t
                            yield t  # 满足要求的k线触发时间
                except Exception as e:
                    print(f"kafka消息错误{e}")
        finally:
            consumer.close()


if __name__ == "__main__":
    symbols = ["BTC/USDT:USDT", "ETH/USDT:USDT", "BTC/USDT", "ETH/USDT"]
    watcher = KafkaWatcherKline("1m", "test_watcher")
    for x in watcher.watch_batch_symbols(symbols):
        print(f"get time {x}")
